# -*- coding: utf-8 -*-
"""\
This file is part of the ptycho package.

Utility function related to experiments.
"""

__all__ = ['raster_scan_positions', 'round_scan_positions', 'round_scan_ROI_positions']

from utils import verbose
import parameters
import numpy as np
import copy
from pyE17 import io
import os

def raster_scan_positions(nx,ny,sx,sy):
    """Raster scan positions (following spec convention)
    
    Parameters
    ----------
    nx, ny : int
        number of scanning steps along x and y (number of point - 1)
    sx, sy : float
        step size along x and y.
    """
    iix, iiy = np.indices((nx+1,ny+1))
    positions = [(sx*i, sy*j) for i,j in zip(iix.ravel(), iiy.ravel())]
    return positions

def round_scan_positions(r_in, r_out, nr, nth):
    """'Round scan' positions (following spec macro convention)
    
    Parameters
    ----------
    r_in : float
        Innermost radius
    r_out : float
        Outermost radius
    nr : int
        number of shell steps ( = number of shells - 1)
    nth : int
        number of points in the inner shell
    """
    dr = (r_out - r_in)/ nr
    positions = []
    for ir in range(1,nr+2):
        rr = r_in + ir*dr
        dth = 2*np.pi / (nth*ir)
        positions.extend([(rr*np.sin(ith*dth), rr*np.cos(ith*dth)) for ith in range(nth*ir)])
    return positions

def round_scan_ROI_positions(dr, lx, ly, nth):
    """'Round scan' positions with ROI (following spec macro convention)
    
    Parameters
    ----------
    dr : float
        shell spacing
    lx, ly : float
        ROI width and height
    nth : int
        number of points in the inner shell    
    """
    rmax = np.sqrt( (lx/2)**2 + (ly/2)**2 )
    nr = np.floor(rmax/dr) + 1
    positions = []
    for ir in range(1,int(nr+2)):
        rr = ir*dr
        th = 2*np.pi*np.arange(nth*ir)/(nth*ir)
        x1 = rr*np.sin(th)
        x2 = rr*np.cos(th)
        positions.extend([(xx1,xx2) for xx1,xx2 in zip(x1,x2) if (np.abs(xx1) <= ly/2) and (np.abs(xx2) <= lx/2)])
    return positions

class DataScan(object):
    """\
    This class represents a single ptychography scan. All parameters are fixed except
    the positions along the two scanning axes.
    """
    def __init__(self, filename=None):
        """\
        Empty object unless a filename is provided
        """
        if filename is None: return
        
        self.load_info(filename)
        self.load_data(filename)

    def load_info(self, filename):
        """\
        Load only information (scan_info) from prepared file.
        """
        self.scan_info = parameters.asParam(io.h5read(filename, 'scan_info')['scan_info'])

    def load_data(self, filename=None, asize=None):
        """\
        Load data.
        """
        if filename is None:
            filename = self.scan_info.data_filename

        if asize is not None:
            if np.isscalar(asize): asize = (asize,asize)
            self.asize = asize
            dpsize = self.scan_info.shape[1:]
            l0,l1 = (dpsize[0]-asize[0])/2., (dpsize[1]-asize[1])/2.
            slice_string = '%d:-%d,%d:-%d' % (int(np.ceil(l0)), int(np.floor(l0)), int(np.ceil(l1)), int(np.floor(l1)))
            self.data = io.h5read(filename, ('data[:,%s]' % slice_string))['data']
            self.fmask = io.h5read(filename, ('fmask[%s]' % slice_string))['fmask']
        else:
            self.asize = self.scan_info.shape[1:]
            tmp = io.h5read(filename,'data', 'fmask')
            self.data = tmp['data']
            self.fmask = tmp['fmask']

    def unload_data(self):
        """\
        Deletes the "data" and "fmask" arrays.
        """
        del self.data
        del self.mask

    def save(self, filename=None, force_overwrite=True):
        """\
        Store this dataset in a standard format.
        If filename is None, use default (scan_info.data_filename)
        """
        # Check if data is available and the right size
        if not hasattr(self, 'data'):
            raise RuntimeError("Attempting to save DataScan instance that does not contain data.")

        if self.data.shape != self.scan_info.shape:
            error_string = "Attempting to save DataScan instance with non-native data dimension "
            error_string += "[data.shape = %s, while scan_info.shape = %s]" % (str(self.data.shape), str(self.scan_info.shape))
            raise RuntimeError(error_string)

        if filename is None:
            filename = self.scan_info.data_filename           

        filename = os.path.expanduser(os.path.abspath(filename))
        write_path = os.path.dirname(filename)
        if not os.path.exists(write_path):
            os.makedirs(write_path)
        if os.path.exists(filename):
            if force_overwrite:
                verbose(1, 'Save file exists but will be overwritten (force_overwrite is True)')
            elif not force_overwrite:
                raise RuntimeError('File %s exists! Operation cancelled.' % filename)
            elif force_overwrite == None:
                ans = raw_input('File %s exists! Overwrite? [Y]/N' % filename)
                if ans and ans.upper() != 'Y':
                    raise RuntimeError('Operation cancelled by user.') 

        h5opt = io.h5options['UNSUPPORTED']
        io.h5options['UNSUPPORTED'] = 'ignore'
        io.h5write(filename, data=self.data, fmask=self.mask, scan_info=self.scan_info._to_dict())
        io.h5options['UNSUPPORTED'] = h5opt
        verbose(3, 'Scan %s data saved to %s.' % (self.scan_info.scan_label, filename))
        return
       
       
class Data(object):
    """\
    Global data container.
    
    Holds all data scans and provides a flat list of all diffraction patterns.
    """
    def __init__(self, scandict):
        pass
        
    def from_scandict(self, scandict):
        """\       
        scandict is a dictionary containing DataScan objects, named by the keys.
        """
        self.initial_scandict = scandict
        
        # Quantities that will have to be uniform for all scans
        wavelengths = []
        detector_distances = []
        detector_pixel_sizes = []
        dpsizes = []
        
        scan_names = []
        scans = {}
        for scan_name, scan in scandict.iteritems():
            scan_names.append(scan_name)
            scans[scan_name] = scan           
            wavelengths.append(scan['wavelength'])
            detector_distances.append(scan['detector_distance'])
            detector_pixel_sizes.append(scan['detector_pixel_size'])
            dpsizes.append(scan['dpsize'])

        # Uniform quantities
        if not np.allclose(wavelengths, wavelengths[0]):
            raise RuntimeError('Reconstruction from scans at different energies is not supported.')
        else:
            self.wavelength = 0. + wavelengths[0]

        if not np.allclose(detector_distances, detector_distances[0]):
            raise RuntimeError('Reconstruction from scans at different detector distances is not supported.')
        else:
            self.detector_distance = 0. + detector_distances[0]

        if not np.allclose(detector_pixel_sizes, detector_pixel_sizes[0]):
            raise RuntimeError('Reconstruction from scans with different detector pixel sizes is not supported.')
        else:
            self.detector_pixel_size = 0. + detector_pixel_sizes[0]

        # Reconstruction pixel size for near-field ptychography
        self.dxnf = self.detector_pixel_size

        if not np.allclose(dpsizes, dpsizes[0]):
            raise RuntimeError('Reconstruction from scans with different array sizes is not supported.')
        else:
            self.dpsize = 0. + dpsizes[0]
                
    def __getitem__(self, key):
        return self.scans[key]
    
    def keys(self):
        """ Same as self.scan_names """
        return self.scans.keys()

    def _prepare_files(self):
        pass
    
    def _load_intensities(self):
        pass        
    def _load_magnitudes(self):
        pass
    
    @property
    def energy(self):
        return 1.2398e-9 / self.wavelength
      
    @property
    def dx(self):
        """\
        Reconstruction pixel size for far-field ptychography.
        """
        return self.wavelength * self.detector_distance / (self.dpsize * self.detector_pixel_size)

